import pandas as pd
import numpy as np
import plotly.graph_objects as go
import gc
from math import pi, sqrt, exp

def plot_by_freq(df, freq, col):
    serie_id = df['series_id'].unique()[0]
    fig = go.Figure()
    # set up trace for wakeup
    df_wakeup = df[df['event']=='wakeup']
    fig.add_traces(go.Ohlc(x=df_wakeup.index,
                    open=df_wakeup[f'open_{col}'], high=df_wakeup[f'high_{col}'],
                    low=df_wakeup[f'low_{col}'], close=df_wakeup[f'close_{col}'],
                   increasing_line_color= 'orange', decreasing_line_color= 'orange', name='awake'
                ))
    # set up traces for onset
    df_onset = df[df['event']=='onset']
    fig.add_traces(go.Ohlc(x=df_onset.index,
                    open=df_onset[f'open_{col}'], high=df_onset[f'high_{col}'],
                    low=df_onset[f'low_{col}'], close=df_onset[f'close_{col}'],
                   increasing_line_color= 'blue', decreasing_line_color= 'blue', name='sleeping'
                ))
    fig.update_layout(title=f'Serie_id: {serie_id} - {freq} OLHC {col}', xaxis_rangeslider_visible=True)
    return fig.to_html(full_html=False)
    
def plot_columns_by_freq(df, freq, columns):
    html_str = ""
    for col in columns:
        df_ohlc = get_df_by_freq(df, freq, col)
        html_str += plot_by_freq(df_ohlc, freq, col)
    return html_str

def load_serie_by_id(idx, df_upload):
    df = df_upload[df_upload['series_id'] == idx]
    df['series_id'] = df['series_id'].astype('category')
    df['timestamp'] = pd.to_datetime(df['timestamp']).apply(lambda t: t.tz_localize(None))
    return df

def get_df_by_freq(df, freq, col):
    df = df.copy()
    df.set_index('timestamp', inplace=True)
    df_ohlc = df[col].resample(freq).ohlc()
    df_ohlc.columns = [f'{c}_{col}' for c in df_ohlc.columns]
    df_ohlc = pd.merge(df_ohlc, df[['event', 'series_id']], left_index=True, right_index=True)
    return df_ohlc

SIGMA = 720 # 12 * 60

def gauss(n=SIGMA,sigma=SIGMA*0.15):
    # guassian distribution function
    r = range(-int(n/2),int(n/2)+1)
    return [1 / (sigma * sqrt(2*pi)) * exp(-float(x)**2/(2*sigma**2)) for x in r]

def get_macd(s, hfast=4, hslow=48, macd=15, center=False):
    sma_fast = s.rolling(hfast*12*60, min_periods=1, center=center).agg('mean')
    sma_slow = s.rolling(hslow*12*60, min_periods=1, center=center).agg('mean')
    macd = (sma_fast - sma_slow).rolling(12*macd, min_periods=1, center=center).mean().astype(np.float32)
    return macd

def macd_feat(df, col='anglezdiffabs', dif=120):
    s = df[f'{col}']
    df[f'{col}_macd'] = get_macd(s)
    s = df[f'{col}'].sort_index(ascending=False)
    df[f'{col}_macd_rev'] = get_macd(s).sort_index()
    df[f'{col}_spred'] = df[f'{col}_macd']-df[f'{col}_macd_rev']
    df[f'{col}_spread_diff']=df[f'{col}_spred'].diff(dif)
    return df

def feat_eng(df):
    df['series_id'] = df['series_id'].astype('category')
    df['timestamp'] = pd.to_datetime(df['timestamp']).apply(lambda t: t.tz_localize(None))
    df['hour'] = df["timestamp"].dt.hour
    df['month'] = df["timestamp"].dt.month
    df.sort_values(['timestamp'], inplace=True)
    df.set_index('timestamp', inplace=True)
    df['enmo'] = df['enmo'].clip(upper=4)
    df['lids'] = np.maximum(0., df['enmo'] - 0.02)
    df['lids'] = df['lids'].rolling(f'{120*5}s', center=True, min_periods=1).agg('sum')
    df['lids'] = 100 / (df['lids'] + 1)
    df['lids'] = df['lids'].rolling(f'{360*5}s', center=True, min_periods=1).agg('mean').astype(np.float32)
    df['lids_sma_fast'] = df['lids'].rolling('2h', center=True, min_periods=1).agg('mean').astype(np.float32)
    df['lids_sma_macd'] = (df['lids'].rolling('18h', center=True, min_periods=1).agg('mean') - df['lids_sma_fast']
                                 ).rolling(f'1h', center=True, min_periods=1).agg('mean').astype(np.float32)
    df["enmo"] = (df["enmo"]*1000).astype(np.int16)
    df["anglez"] = df["anglez"].astype(np.int16)
    df["anglezdiffabs"] = df["anglez"].diff().abs().astype(np.float32)
    for col in ['enmo', 'anglez', 'anglezdiffabs']:
        periods = [720] 
        for n in periods:
            rol_args = {'window':f'{n+5}s', 'min_periods':10, 'center':True}
            for agg in ['std', 'mean']:
                df[f'{col}_{agg}_{n}'] = df[col].rolling(**rol_args).agg(agg).astype(np.float32).values
                gc.collect()
            rol_args = {'window':f'8h', 'min_periods':15, 'center':True}
            df[f'{col}_std_max_8h'] = df[f'{col}_std_{n}'].rolling(**rol_args).max().astype(np.float32).values
    df.reset_index(inplace=True)
    return df

def plot_submission(submission, df_upload) :
  test_events = submission.drop(['row_id', 'score'], axis=1)
  test_events['timestamp'] = pd.to_datetime(test_events['timestamp']).apply(lambda t: t.tz_localize(None))
  count_na = test_events.groupby('series_id')['timestamp'].apply(lambda x: x.isna().sum())
  html_str = ""
  for idx in count_na[count_na>=0].index[:3]:
    print(f'ID: {idx}', "="*120, sep='\n')
    test_series = load_serie_by_id(idx, df_upload)
    df = test_series.merge(test_events[test_events['series_id']==idx], on=['series_id', 'timestamp'], how='left').ffill()
    df = feat_eng(df)
    html_str += plot_columns_by_freq(df, '5min', ['anglez', 'anglezdiffabs_mean_720', 'enmo', 'enmo_mean_720'])
  return html_str
